/**
 * @file AtpgLibrary.cc
 * @author Xiaoze Lin (linxiaoze96@gmail.com)
 * @brief The implementation of the Atpg library parser.
 * @version 0.1
 *
 * @copyright Copyright (c) 2024
 *
 */
#include "AtpgLibrary.hh"

namespace ifsim {

/**
 * @brief: parser the standard cell library file.
 * @param {char*} library_name
 * @return none
 */
void AtpgLibrary::parserStdCells(const char* library_name)
{
  std::ifstream ifs(library_name);
  if (!ifs) {
    LOG_ERROR << "Failed to read json file: '" << library_name << "'.";
  }
  Json* json = new Json();
  ifs >> *json;

  jsonToStdCells(json);

  ifs.close();
  delete json;
}

/**
 * @brief: convert json to standard cells information.
 * @param {Json*} json
 * @return none
 */
void AtpgLibrary::jsonToStdCells(Json* json)
{
  auto std_cells = json->at("std_cells");
  int cell_index = 0;
  for (auto& std_cell_iter : std_cells) {
    StdCellModel cell;
    unsigned node_id = 0;
    std::map<std::string, unsigned> net;  // <net name, node_id>

    // cell name
    auto name = std_cell_iter.at("name");
    for (auto& name_iter : name) {
      _nameMap.insert(std::pair<std::string, int>(name_iter, cell_index));
      cell.cellName(name_iter);
    }
    ++cell_index;

    // cell signals
    auto signal = std_cell_iter.at("signals");
    for (auto& signal_iter : signal) {
      // record inputs
      auto input = signal_iter.at("input");
      cell.cellInputNum(input.size());
      for (auto& input_iter : input) {
        cell.cellInputs(input_iter);
        cell.cellInputNodes(node_id);
        net.insert(std::pair<std::string, unsigned>(input_iter, node_id++));
      }

      // record wires
      auto wire = signal_iter.at("wire");
      for (auto& wire_iter : wire) {
        net.insert(std::pair<std::string, unsigned>(wire_iter, node_id++));
      }

      // record outputs finally, make sure outputs are in the end of map "net"
      auto output = signal_iter.at("output");
      int output_node_id = (node_id == 0) ? output.size() : node_id;
      cell.cellOutputNum(output.size());
      for (auto& output_iter : output) {
        cell.cellOutputs(output_iter);
        cell.cellOutputNodes(output_node_id++);
        net.insert(std::pair<std::string, unsigned>(output_iter, node_id++));
      }

      // build primitive nodes in std cell model, include inputs, pri nodes, outputs
      cell.initPrimitiveNode(node_id + cell.cellOutputNum());  // num of nodes + num of outputs
      // record input nodes
      for (int i = 0; i < cell.cellInputNum(); ++i) {
        PrimitiveGate inputNode;
        inputNode.gid(i);
        inputNode.type(T_INPUT);  // not truely PI
        cell.primitiveNode(inputNode);
      }
      // record primitives
      auto primitive = std_cell_iter.at("sim_primitives");
      cell.primitiveNum(primitive.size());
      for (auto& primitive_iter : primitive) {
        auto sim_type = primitive_iter.at("sim_type");
        auto connection = primitive_iter.at("connection");
        PrimitiveGate prim_gate;
        prim_gate.type(LUT_StringToGType[sim_type]);
        std::string net_name = connection[0];
        LOG_FATAL_IF(net.find(net_name) == net.end()) << "Cannot find the net name " << net_name << " in " << cell.cellName()[0];
        prim_gate.gid(net[net_name]);

        // look for pred primitive gates (if any)
        std::vector<int> pred_nodes;
        for (int i = 1; i < (int) connection.size(); ++i) {
          net_name = connection[i];
          if (("1'b0" != net_name) && ("1'b1" != net_name) && (net.find(net_name) == net.end())) {
            LOG_FATAL_IF(net.find(net_name) == net.end()) << "Cannot find the net name " << net_name << " in " << cell.cellName()[0];
          } else if ("1'b0" == net_name) {
            pred_nodes.push_back(tie0);
          } else if ("1'b1" == net_name) {
            pred_nodes.push_back(tie1);
          } else {
            pred_nodes.push_back(net[net_name]);
          }
        }
        prim_gate.preds(pred_nodes);
        cell.primitiveNode(prim_gate);
      }
      // record output nodes
      int output_node_start_id = node_id;
      for (int i = 0; i < cell.cellOutputNum(); ++i) {
        PrimitiveGate outputNode;
        outputNode.gid(output_node_start_id + i);
        outputNode.type(T_OUTPUT);  // not truely PO

        // look for pred primitive node
        std::vector<int> pred_nodes;
        std::string output_net_name = cell.cellOutputs(i);
        pred_nodes.push_back(net[output_net_name]);
        outputNode.preds(pred_nodes);

        cell.primitiveNode(outputNode);
      }

      // look for succ primitive gates (if any)
      for (int i = cell.wholePrimitiveNum() - 1; i >= 0; --i) {
        const PrimitiveGate& tmp_prim_gate = cell.primitiveNode(i);
        for (auto pred_id : tmp_prim_gate.preds()) {
          if ((pred_id == tie0) || (pred_id == tie1)) {
            continue;
          } else {
            PrimitiveGate tmp_pred_gate = cell.primitiveNode(pred_id);
            std::vector<int> tmp_succs = tmp_pred_gate.succs();
            tmp_succs.push_back(i);
            tmp_pred_gate.succs(tmp_succs);
            cell.primitiveNode(tmp_pred_gate);
          }
        }
      }

      insertBufferAfterInputFanoutStem(cell);
      insertBufferAfterOutputFanoutStem(cell);

      _stdCells.push_back(cell);
    }
  }
}

/**
 * @brief: Insert buffer after input fanout stem node to make sure one input map to one pri node.
 */
void AtpgLibrary::insertBufferAfterInputFanoutStem(StdCellModel& cell)
{
  auto pri_nodes = cell.primitiveNodes();
  std::map<int, int> pi_to_num_add_buf;
  int input_num = cell.cellInputNum();
  int num_add_buf = 0;

  // count num of buffer need to be added
  for (int pid = 0; pid < input_num; ++pid) {
    auto pri_node = pri_nodes[pid];
    if (pri_node.numSuccs() == 1) {
      continue;
    }
    pi_to_num_add_buf.insert(std::pair<int, int>(pid, num_add_buf));
    ++num_add_buf;
  }

  if (num_add_buf > 0) {
    int num_new_pnode = input_num;
    std::vector<PrimitiveGate> tmp_pri_nodes(pri_nodes.size() + num_add_buf);
    // input pri node
    for (int pid = 0; pid < input_num; ++pid) {
      PrimitiveGate pri_node = pri_nodes[pid];
      if (pi_to_num_add_buf.find(pid) != pi_to_num_add_buf.end()) {  // input fanout stem
        std::vector<int> tmp_succs(1, input_num + pi_to_num_add_buf[pid]);
        pri_node.succs(tmp_succs);  // update succs of pi

      } else {  // input pri node withou fanout
        std::vector<int> tmp_succs(1, pri_node.succ(0) + num_add_buf);
        pri_node.succs(tmp_succs);  // update succs of pi
      }
      tmp_pri_nodes[pri_node.gid()] = pri_node;
    }
    // added buffer (add between input pri nodes and other pri nodes)
    for (auto add_buf : pi_to_num_add_buf) {
      PrimitiveGate pi_pnode = pri_nodes[add_buf.first];
      int new_gid = num_new_pnode++;
      PrimitiveGate buf_pnode;
      buf_pnode.gid(new_gid);
      buf_pnode.type(T_BUF);
      std::vector<int> tmp_preds(1, pi_pnode.gid());
      buf_pnode.preds(tmp_preds);  // update preds of added buffer
      std::vector<int> tmp_succs;
      for (auto succ_id : pi_pnode.succs()) {
        tmp_succs.push_back(succ_id + num_add_buf);
      }
      buf_pnode.succs(tmp_succs);  // update succs of added buffer
      tmp_pri_nodes[new_gid] = buf_pnode;
    }
    // other pri nodes and output pri nodes
    for (int pid = input_num; pid < pri_nodes.size(); ++pid) {
      PrimitiveGate pri_node = pri_nodes[pid];
      int new_gid = num_new_pnode++;
      pri_node.gid(new_gid);  // update gid
      std::vector<int> tmp_preds(pri_node.numPreds());
      for (int pred_i = 0; pred_i < pri_node.numPreds(); ++pred_i) {
        int pred_id = pri_node.pred(pred_i);
        if (pi_to_num_add_buf.find(pred_id) != pi_to_num_add_buf.end()) {  // pred is input fanout stem
          tmp_preds[pred_i] = input_num + pi_to_num_add_buf[pred_id];

        } else if (pred_id >= input_num) {  // pred is other pri nodes
          tmp_preds[pred_i] = pred_id + num_add_buf;

        } else {  // pred is input pri node but not fanout stem
          tmp_preds[pred_i] = pred_id;
        }
      }
      pri_node.preds(tmp_preds);  // update preds
      std::vector<int> tmp_succs(pri_node.numSuccs());
      for (int succ_i = 0; succ_i < pri_node.numSuccs(); ++succ_i) {
        int succ_id = pri_node.succ(succ_i);
        tmp_succs[succ_i] = succ_id + num_add_buf;  // succ must after all added buffers in new pri nodes
      }
      pri_node.succs(tmp_succs);  // update succs
      tmp_pri_nodes[new_gid] = pri_node;
    }
    cell.primitiveNodes(tmp_pri_nodes);  // update pri nodes in std cell model
    cell.primitiveNum(cell.primitiveNum() + num_add_buf);
  }
}

/**
 * @brief: Insert buffer after output fanout stem node to make sure one output map to one pri node.
 */
void AtpgLibrary::insertBufferAfterOutputFanoutStem(StdCellModel& cell)
{
  auto pri_nodes = cell.primitiveNodes();
  std::map<int, int> po_to_num_add_buf;
  int input_num = cell.cellInputNum();
  int output_num = cell.cellOutputNum();
  int pri_num = cell.primitiveNum();
  int output_start_id = input_num + pri_num;
  int num_add_buffer = 0;

  // count num of buffer need to added
  for (int output_i = 0; output_i < output_num; ++output_i) {
    int output_id = output_i + output_start_id;
    PrimitiveGate output_pnode = pri_nodes[output_id];
    if (pri_nodes[output_pnode.pred(0)].numSuccs() > 1) {
      po_to_num_add_buf.insert(std::pair<int, int>(output_id, num_add_buffer));
      ++num_add_buffer;
    }
  }

  if (num_add_buffer > 0) {
    int num_new_pnode = output_start_id;
    std::vector<PrimitiveGate> tmp_pri_nodes(pri_nodes.size() + num_add_buffer);
    // input pri node
    for (int pid = 0; pid < input_num; ++pid) {
      PrimitiveGate pri_node = pri_nodes[pid];
      tmp_pri_nodes[pri_node.gid()] = pri_node;
    }
    // other pri nodes except input and output pri nodes
    for (int pid = input_num; pid < output_start_id; ++pid) {
      PrimitiveGate pri_node = pri_nodes[pid];
      for (auto& succ_id : pri_node.updateSuccs()) {
        if (po_to_num_add_buf.find(succ_id) != po_to_num_add_buf.end()) {  // succ is output fanout stem
          succ_id = output_start_id + po_to_num_add_buf[succ_id];          // update succs
        } else if (succ_id >= output_start_id) {                           // succ is output pri node but not fanout stem
          succ_id = succ_id + num_add_buffer;                              // update succs
        }
      }
      tmp_pri_nodes[pri_node.gid()] = pri_node;
    }
    // added buffer (after other pri nodes, before output pri nodes)
    for (auto add_buf : po_to_num_add_buf) {
      int add_buf_i = add_buf.second;
      PrimitiveGate po_pnode = pri_nodes[add_buf.first];
      int new_gid = num_new_pnode++;
      PrimitiveGate buf_pnode;
      buf_pnode.gid(new_gid);
      buf_pnode.type(T_BUF);
      std::vector<int> tmp_preds(1, po_pnode.pred(0));
      buf_pnode.preds(tmp_preds);  // update preds of added buffer
      std::vector<int> tmp_succs(1, po_pnode.gid() + num_add_buffer);
      buf_pnode.succs(tmp_succs);  // update succs of added buffer
      tmp_pri_nodes[new_gid] = buf_pnode;
    }
    // output pri node
    for (int output_i = 0; output_i < output_num; ++output_i) {
      int output_id = output_i + output_start_id;
      PrimitiveGate po_pnode = pri_nodes[output_id];
      if (po_to_num_add_buf.find(output_id) != po_to_num_add_buf.end()) {
        std::vector<int> tmp_preds(1, output_start_id + po_to_num_add_buf[output_id]);
        po_pnode.preds(tmp_preds);  // update preds
      }
      int new_gid = num_new_pnode++;
      po_pnode.gid(new_gid);
      tmp_pri_nodes[new_gid] = po_pnode;
    }
    cell.primitiveNodes(tmp_pri_nodes);  // update pri nodes in std cell model
    cell.primitiveNum(cell.primitiveNum() + num_add_buffer);
  }
}

/**
 * @brief: Find cell according to cell name.
 */
StdCellModel* AtpgLibrary::findCell(std::string cname)
{
  if (_nameMap.count(cname)) {
    return &_stdCells[_nameMap[cname]];
  } else {
    LOG_ERROR << "there is no standard cell named '" << cname << "'";
    return NULL;
  }
}
}  // namespace ifsim